'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2024-09-24
Copyright © Department of Physics, Tsinghua University. All rights reserved

Model for Zhipu AI
'''

from .common import AIBase
from zhipuai import ZhipuAI


class ZhipuLLM(AIBase):
    client: ZhipuAI

    def __init__(self, user_key: str, model_name: str, url: str) -> None:
        self.client = ZhipuAI(api_key=user_key)
        super().__init__(user_key, model_name, url)

    def get_content(self, messages: list[dict[str, str]]) -> str:
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            top_p=0.7,
            temperature=0.95,
            max_tokens=1024,
        )
        return response.choices[0].message.content

    def batch_run(self, message_list: list[list[dict]], source_file: str, target_file: str):
        ''' Current implementation: run 'get_content' for many times '''
        return [self.get_content(messages) for messages in message_list]
